// Dll_Injector.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <windows.h>
#include <string>
#include <TlHelp32.h>

namespace std
{
#ifdef UNICODE
	typedef wstring tstring;
#else
	typedef string tstring;
#endif
}

bool InjectDll(DWORD processID, const std::tstring& dllFilePath)
{
#ifdef UNICODE
	const char* pApiName = "LoadLibraryW";
#else
	const char* pApiName = "LoadLibraryA";
#endif

	HANDLE hProcess = OpenProcess(
							PROCESS_QUERY_INFORMATION|
							PROCESS_CREATE_THREAD|
							PROCESS_VM_OPERATION|
							PROCESS_VM_WRITE
						,	FALSE
						,	processID);

	size_t size_of_dllFilePath = (dllFilePath.size()+1)*sizeof(TCHAR);

	char* pDllFilePathForRemote = 
		(char*)VirtualAllocEx(
							hProcess
						,	nullptr
						,	size_of_dllFilePath
						,	MEM_COMMIT
						,	PAGE_READWRITE);

	WriteProcessMemory(
		hProcess
		,	pDllFilePathForRemote
		,	dllFilePath.c_str()
		,	size_of_dllFilePath
		,	nullptr);


	PTHREAD_START_ROUTINE pfnThread = 
		(PTHREAD_START_ROUTINE)GetProcAddress(
					GetModuleHandle(_T("kernel32"))
				,	pApiName);

	HANDLE hRemoteThread = CreateRemoteThread(
									hProcess
								,	nullptr
								,	0
								,	pfnThread
								,	pDllFilePathForRemote
								,	0
								,	nullptr);

	WaitForSingleObject(hRemoteThread, INFINITE);

	VirtualFreeEx(
			hProcess
		,	pDllFilePathForRemote
		,	0
		,	MEM_RELEASE);
	
	CloseHandle(hRemoteThread);
	CloseHandle(hProcess);
	return true;
}

bool EjectDll(DWORD processID, const std::tstring& dllFilePath)
{

	const char* pApiName = "FreeLibrary";

	HANDLE hSnapshot = CreateToolhelp32Snapshot(
							TH32CS_SNAPMODULE
						,	processID);

	if (hSnapshot == INVALID_HANDLE_VALUE)
	{
		return false;
	}
	
	MODULEENTRY32 entry;
	entry.dwSize = sizeof(entry);

	BOOL existModule = Module32First(hSnapshot, &entry);
	for(; existModule; existModule = Module32Next(hSnapshot, &entry))
	{
		BOOL isFound = (dllFilePath == entry.szModule)
					|| (dllFilePath == entry.szExePath);

		if (isFound)
		{
			HANDLE hProcess = OpenProcess(
									PROCESS_QUERY_INFORMATION|
									PROCESS_CREATE_THREAD|
									PROCESS_VM_OPERATION					
								,	FALSE
								,	processID);

			PTHREAD_START_ROUTINE pfnThread = 
				(PTHREAD_START_ROUTINE)GetProcAddress(
					GetModuleHandle(_T("kernel32"))
				,	pApiName);

			HANDLE hRemoteThread = CreateRemoteThread(
										hProcess
									,	nullptr
									,	0
									,	pfnThread
									,	entry.modBaseAddr
									,	0
									,	nullptr);

			WaitForSingleObject(hRemoteThread, INFINITE);
			
			CloseHandle(hRemoteThread);
			CloseHandle(hProcess);
			CloseHandle(hSnapshot);
			return true;
		}
	}

	return false;
}

int _tmain(int argc, _TCHAR* argv[])
{
	InjectDll(_ttoi(argv[1]), argv[2]);
	//EjectDll(_ttoi(argv[1]), argv[2]);
	return 0;
}

